import datetime
import json
import os
import threading
import time
from queue import Queue

import av
import cv2
import numpy as np
import requests
from ultralytics import YOLO

from logger import log
from settings import settings

PATH = os.path.abspath('.') + '/img_zp/'
if not os.path.exists(PATH):
    print("文件夹没有，创建img_zp文件夹")
    os.makedirs(PATH)  # 如果不存在则创建目录

# model = YOLO("yolov8n.pt")
model = YOLO("best_people_20250729.pt")

model.info()


def post_img(frame, stream_id):
    # 将图像编码为JPEG格式
    _, buffer = cv2.imencode('.jpg', frame)
    # 转换为字节
    jpg_as_text = buffer.tobytes()
    try:
        # 创建表单数据
        files = {'contractFile': ('img.jpg', jpg_as_text)}
        # 设置请求的URL
        url = 'https://' + settings.host + '/api/account/oss/fileUpload?accessCode=' + settings.code
        # 发送POST请求
        response = requests.post(url, files=files, verify=False, timeout=10)
        response.raise_for_status()  # 如果请求返回了错误状态码，这会抛出异常
        # 处理响应
        # log.info(f"图片上传数据返回:{response.json()}")
        snapshotUrl = response.json()['data']
        params = {
            "id": stream_id,
            "snapshotUrl": snapshotUrl,
            "platformId": settings.platformId,
            "platformName": settings.platformName,
            "type": "invade"
        }
        requests_post = requests.post('https://' + settings.host + '/api/weigh/internal/vs_stream/snapshot',
                                      json=params)
        requests_post.raise_for_status()  # 如果请求返回了错误状态码，这会抛出异常
        # log.info(f"物体检测在区域内，上传完成:{params}:{requests_post.json()}")
    except Exception as error:
        # 打印错误信息
        log.info("post_img error:", error)
    # cv2.imwrite(f'capture/capture_{timestamp}.jpg', save_image)


def is_inside_polygon(points, polygon):
    """ 检查点是否在多边形内 """
    polygon_points = np.array(polygon, np.int32)
    polygon_points = polygon_points.reshape((-1, 1, 2))

    for point in points:
        # 转换点为元组 (x, y)
        pt = (int(point[0]), int(point[1]))
        if cv2.pointPolygonTest(polygon_points, pt, False) >= 0:
            return True
    return False


def draw_polygon_points(image, polygon_points):
    # 绘制检测框
    # [[400, 300], [400, 500], [600, 500], [600, 300]]
    # 定义颜色和线条粗细
    color = (0, 255, 0)  # 绿色
    thickness = 3
    # 创建一个表示多边形顶点的数组
    points = np.array(polygon_points, np.int32)
    points = points.reshape((-1, 1, 2))
    # 绘制多边形
    cv2.polylines(image, [points], True, color, thickness)


class ObjectDetection:
    def __init__(self):
        self.frame_queue = Queue()

    def screenshot(self, rtsp_url):
        cap = None
        try:
            cap = cv2.VideoCapture(rtsp_url)
            if cap.isOpened():
                # 寻找关键帧
                for i in range(25):
                    cap.read()
                ret, frame = cap.read()
                if ret:
                    return frame
                else:
                    # log.warn(f"ret: {ret}")
                    return None
            else:
                # log.warn(f"cap.isOpened() False")
                return None
        except Exception as e:
            log.exception(f"e:{e}")
        finally:
            # 释放资源
            cap.release()
        return None

    def screenshot_av(self, rtsp_url):
        try:
            container = av.open(rtsp_url,
                                format='rtsp',
                                options={'rtsp_transport': 'tcp', 'video_size': '640x360'},
                                timeout=25)
        except Exception as e:
            # log.exception(e)
            # log.warn(f"Failed to open RTSP stream: {e}")
            return None
        for frame in container.decode(video=0):
            if frame.key_frame:
                # 将帧转换为numpy数组
                img = frame.to_ndarray(format='bgr24')
                return img

    def thread_ob(self, i):
        while True:
            try:
                item = settings.streams[i]
                item_id = item['id']
                item_name = item['name']
                rtsp_url = item['rtsp']
                if item['cameraRtsp'] is not None and item['cameraRtsp'] != '':
                    rtsp_url = item['cameraRtsp']
                polygon_points_str = item['polygonPoints']
                if polygon_points_str != '' and polygon_points_str != '[]':
                    # log.info(f"{i}|{item_id}|{item_name}|开始取流")
                    polygon_points = json.loads(polygon_points_str)
                    frame = self.screenshot_av(rtsp_url)
                    # frame = self.screenshot(rtsp_url)
                    if frame is not None:
                        log.info(f"{i}|{item_id}|{item_name}|取流成功")
                        frame = cv2.resize(frame, (640, 360))
                        # 复制frame，保存一个不绘制的帧
                        frame_without_polygon = frame.copy()
                        # 绘制检测框
                        draw_polygon_points(frame, polygon_points)
                        # 对获取的帧进行处理
                        results = model.predict(frame, verbose=False, classes=[0, ], conf=0.7)
                        annotated_frame = results[0].plot()

                        if len(results[0].boxes) > 0:
                            for box in results[0].boxes:
                                class_name = model.names[int(box.cls)]
                                log.info(
                                    f"{i}|{item_id}|{item_name}|检测到物体: {class_name} (置信度: {float(box.conf[0]):.1f})")
                                x1, y1, x2, y2 = box.xyxy[0]
                                bbox_points = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
                                # 检查物体是否在多边形内
                                if is_inside_polygon(bbox_points, polygon_points):
                                    # 保存本地，为了训练模型用；
                                    now = datetime.datetime.now()
                                    timestamp = now.strftime('%Y-%m-%d-%H-%M-%S')
                                    image_name = f"img_zp/{timestamp}_{item_id}.jpg"
                                    cv2.imwrite(image_name, frame_without_polygon)
                                    # 物体在区域内，添加时间戳并保存图像
                                    # log.info(f"{i}|{item_id}|{item_name}|thread_ob: 物体检测在区域内，开始上传")
                                    post_img(annotated_frame, item_id)

                        # 获取检测到的物体的边界框坐标
                        self.frame_queue.put(annotated_frame)
                    else:
                        log.info(f"{i}|{item_id}|{item_name}|取流失败|{rtsp_url}")
            except Exception as e:
                log.exception(e)
            time.sleep(5)

    def thread_loop(self):
        while True:
            frame = self.frame_queue.get()
            cv2.imshow(f"frame", frame)
            cv2.waitKey(1)

    def start(self):
        for i in range(len(settings.streams)):
            threading.Thread(target=self.thread_ob, args=(i,), daemon=True).start()

        threading.Thread(target=self.thread_loop, daemon=True).start()


objectdetection = ObjectDetection()

if __name__ == '__main__':
    # polygon_points = np.array([[400, 300], [400, 500], [600, 500], [600, 300]], np.int32)
    # polygon_points = polygon_points.reshape((-1, 1, 2))
    # cap = cv2.VideoCapture('rtsp://admin:admin@192.168.110.100:8557/h264')
    settings.streams = [
        {
            'id': 1,
            'name': "name",
            'rtsp': 0,
            'polygonPoints': '[[100, 100], [100, 1800], [1800, 1800], [1800, 100]]'
        }
    ]
    objectdetection.start()
